import uuid
import json
import itertools
import os
import requests
import aiofiles
import aiofiles.os

import tornado
from tornado.websocket import WebSocketHandler
from tornado.web import RequestHandler

from log import Log
from models.client import InteractiveClient
from models import auths
import config


class CompleteHandler(RequestHandler):
    def set_default_headers(self):
        self.set_header("Access-Control-Allow-Origin", "*")  # 这个地方可以写域名
        self.set_header("Access-Control-Allow-Headers", "x-requested-with")
        self.set_header("Access-Control-Allow-Methods", "POST, GET, OPTIONS")

    def options(self):
        # no body
        self.set_status(204)
        self.finish()

    def get(self, *args, **kwargs):
        code = self.get_argument("code", None)
        kernel = self.application.kernel["kernel"]
        if code is not None or kernel is not None:
            results = []
            for state in itertools.count():
                comp = kernel.complete(code, state)
                if comp is None:
                    break
                results.append(comp)

            self.write({
                "is_success": True,
                "code": "",
                "data": {"code": code, "content": results}
            })
        else:
            pass


class WeChatHandler(WebSocketHandler):
    def __init__(self, application, request):
        super().__init__(application, request)
        self.logger = Log().logger
        self.connected = False

    def check_origin(self, origin):
        return True

    def get_current_user(self):
        token = self.get_argument(name="token", default=None)
        if token is None:
            return False
        user_id = auths.get_user_id(token)
        if user_id is None:
            return False
        return True

    @tornado.web.authenticated
    async def open(self):
        self.connected = True
        if self.application.kernel is None:
            kernel = InteractiveClient(self.application.kernel_ip, self.application.kernel_port)
            kernel.register("input", self._handle_input)
            kernel.register("output", self._handle_output)
            kernel.register("err", self._handle_err)
            self.application.kernel = {"client": self,
                                       "kernel": kernel,
                                       "source_mode": "single",
                                       "output_cached": []}
        msg_id = "id_" + str(uuid.uuid4())
        message = {"id": msg_id,
                   "type": "output",
                   "data": {"type": "text", "content": "登录了"},
                   "desc": {}}
        await self.recv(message)

    async def on_message(self, message):
        self.logger.info(f"recv message: {message}")
        msg_code_id = "id_" + str(uuid.uuid4())
        if self.application.kernel is not None:
            kernel = self.application.kernel["kernel"]
            if kernel is not None:
                msg_obj = json.loads(message)
                # edit mode
                if msg_obj["mode"] == "edit":
                    if msg_obj["language"] == "python3":
                        msg_obj_data = msg_obj["data"]
                        await self.recv({"id": msg_code_id, "type": "code", "data": msg_obj_data, "desc": {}})
                        msg_obj_code = ""
                        if isinstance(msg_obj_data, str):
                            msg_obj_code = msg_obj_data.strip()
                        elif isinstance(msg_obj_data, dict) and "path" in msg_obj_data:
                            msg_obj_path = msg_obj_data["path"].strip()
                            async with aiofiles.open(msg_obj_path, 'r', encoding='utf-8', errors='ignore') as f:
                                msg_obj_code = await f.read()
                                msg_obj_code = msg_obj_code.strip()
                        else:
                            pass
                        if not msg_obj_code == "":
                            if "option" in msg_obj and "source_mode" in msg_obj["option"] \
                                    and msg_obj["option"]["source_mode"] == "exec":
                                self.application.kernel["source_mode"] = "exec"
                                self.application.kernel["output_cached"] = []
                                msg_obj_code = '\n'.join(['    ' + line for line in msg_obj_code.split("\n")])
                                msg_obj_code = "try:\n" + msg_obj_code + "\nfinally:\n    print(\"*Close*\")"
                            else:
                                self.application.kernel["source_mode"] = "single"
                            msg_obj_desc = getattr(msg_obj, "desc", "")
                            if "option" in msg_obj and "debug" in msg_obj["option"] \
                                    and msg_obj["option"]["debug"] is True:
                                current_sequence = await self._debug_python3(kernel, msg_obj["type"],
                                                                             msg_code_id,
                                                                             msg_obj_code,
                                                                             msg_obj_desc)
                            else:
                                current_sequence = await self._handle_python3(kernel, msg_obj["type"],
                                                                              msg_code_id,
                                                                              msg_obj_code,
                                                                              msg_obj_desc)
                            await self.recv({"id": msg_code_id, "type": "sequence", "data": current_sequence, "desc": {}})
                        else:
                            pass
                    elif msg_obj["language"] == "markdown":
                        pass
                    else:
                        pass
                # command mode
                elif msg_obj["mode"] == "command":
                    if msg_obj["type"] == "heartbeat":
                        await self.recv({"id": msg_code_id, "type": "heartbeat"})
                    else:
                        pass
                else:
                    pass
            else:
                pass
        else:
            pass

    async def recv(self, msg):
        try:
            if self.connected:
                client = self.application.kernel["client"]
                await client.write_message(msg)
            else:
                pass
        except BaseException as e:
            self.logger.error(f"{e}", exc_info=True)

    def on_close(self):
        self.connected = False
        requests.get("http://localhost:8090/close?port=%d&kernel_ip=%s&kernel_port=%d" % (self.application.server_port,
                                                                                          self.application.kernel_ip,
                                                                                          self.application.kernel_port))
        if self.application.kernel is not None:
            del self.application.kernel
        else:
            pass

    async def _debug_python3(self, kernel, msg_obj_type, msg_code_id, msg_obj_code, msg_obj_desc):
        if msg_obj_type == "compute":
            return await kernel.runcode(msg_code_id, msg_obj_code, msg_obj_desc, True)
        else:
            pass
        return 0

    async def _handle_python3(self, kernel, msg_obj_type, msg_code_id, msg_obj_code, msg_obj_desc):
        if msg_obj_type == "compute":
            return await kernel.runcode(msg_code_id, msg_obj_code, msg_obj_desc, False)
        else:
            pass
        return 0

    async def _handle_input(self, pid, code, desc):
        await self.recv({"id": pid, "type": "input", "code": code, "desc": desc})

    async def _handle_output(self, omsg, pid, code, desc):
        handler = getattr(self, "handle_%s" % omsg.msg_type)
        if handler is not None:
            await handler(omsg, pid, code, desc)

    async def _handle_err(self, err, pid, code, desc):
        self.logger.error(f"{err.etype} : {err.evalue}", exc_info=False)
        self.logger.error(f"{''.join(err.traceback)}", exc_info=False)
        msg_id = "id_" + str(uuid.uuid4())
        message = {"id": msg_id, "type": "error", "code": code, "data": err.traceback, "desc": desc}
        await self.recv(message)

    async def handle_pyin(self, omsg, pid, code, desc):
        pass

    async def handle_pyout(self, omsg, pid, code, desc):
        msg_id = "id_" + str(uuid.uuid4())
        category = getattr(omsg.content.data, "category", None)
        content = getattr(omsg.content.data, "content", None)
        if category is None:
            resp = content.rstrip()
            if not resp == "":
                message = {"id": msg_id, "type": "output", "code": code,
                           "data": {"type": "text", "content": resp},
                           "desc": desc}
                if self.application.kernel["source_mode"] == "exec":
                    self.application.kernel["output_cached"].append(message)
                else:
                    await self.recv(message)
            else:
                pass
        else:
            message = {"id": msg_id, "type": "output", "code": code,
                       "data": {"type": category, "content": content},
                       "desc": desc}
            if self.application.kernel["source_mode"] == "exec":
                self.application.kernel["output_cached"].append(message)
            else:
                await self.recv(message)

    async def handle_stream(self, omsg, pid, code, desc):
        if omsg.content.name == "stdout":
            resp = omsg.content.data.rstrip()
            if not resp == "":
                # 进度条
                if resp.startswith('\x08'):
                    await self.recv({"id": pid, "type": "progress", "data": resp.strip('\x08'), "desc": desc})
                else:
                    msg_id = "id_" + str(uuid.uuid4())
                    message = {"id": msg_id, "type": "stream", "code": code, "data": resp, "desc": desc}
                    if self.application.kernel["source_mode"] == "exec":
                        if resp == "*Close*":
                            if config.options["rust"]:
                                file_path = os.path.join(os.getcwd(), "output")
                                if not os.path.exists(file_path):
                                    await aiofiles.os.mkdir(file_path)
                                else:
                                    pass
                                async with aiofiles.open(os.path.join(file_path, msg_id), mode="w+") as f:
                                    for line in self.application.kernel["output_cached"]:
                                        await f.write(json.dumps(line) + "\n")
                            else:
                                pass
                            await self.recv(json.dumps(self.application.kernel["output_cached"]))
                            self.application.kernel["output_cached"] = []
                        else:
                            self.application.kernel["output_cached"].append(message)
                    else:
                        await self.recv(message)
            else:
                pass
        else:
            self.logger.error("*ERR*", exc_info=False)
            self.logger.error(f"{omsg.content.data}", exc_info=False)

    async def handle_pyerr(self, omsg, pid, code, desc):
        pass


class InputHandler(WebSocketHandler):
    def __init__(self, application, request):
        super().__init__(application, request)

    def check_origin(self, origin):
        return True

    def get_current_user(self):
        token = self.get_argument(name="token", default=None)
        if token is None:
            return False
        user_id = auths.get_user_id(token)
        if user_id is None:
            return False
        return True

    @tornado.web.authenticated
    def open(self):
        pass

    def on_message(self, message):
        kernel = self.application.kernel["kernel"]
        if kernel is not None:
            kernel.input(message)
